"""
@Author    : ghenyar
@Time      : 2025/9/30 16:39
@File      : repository
@Desc      : SQLAlchemy 二次封装
"""
from typing import Type, TypeVar, Optional, Sequence, Any, List, Callable, Union
from sqlalchemy import select, func, update as updateData, exists as existsData, delete as deleteData
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.orm import load_only

from app.common.page_result import PageResult

ModelT = TypeVar("ModelT")

ColumnLike = Union[ColumnElement, Any]


class ModelQuerier:
    def __init__(self, session: AsyncSession, model: Type[ModelT]):
        self.session = session  # 数据库会话
        self.model = model  # 模型
        self._where: List[ColumnElement] = []  # 查询条件
        self._order_by: list[ColumnElement] = []  # 排序
        self._cols: Sequence[ColumnLike] | None = None

    def where(self, *where: ColumnElement) -> "ModelQuerier":
        """查询条件"""
        self._where.extend(where)
        return self

    def order_by(self, *order: ColumnElement) -> "ModelQuerier":
        """排序"""
        self._order_by.extend(order)
        return self

    def only(self, *cols: ColumnLike):
        self._cols = cols
        return self

    async def update(self, **values) -> int | Callable[[], int]:
        """
        更新数据
        :param values: 要更新的字段和值
        :return: 更新的行数
        """
        if not values:
            return 0

        stmt = updateData(self.model).where(*self._where).values(**values)
        result = await self.session.execute(stmt)
        await self.session.commit()
        return result.rowcount

    async def update_one(self, **values) -> bool:
        """
        更新单条数据
        :param values: 要更新的字段和值
        :return: 是否更新成功
        """
        rowcount = await self.update(**values)
        return rowcount > 0

    async def exists(self) -> bool:
        """
        检查是否存在满足条件的数据
        :return: True/False
        """
        stmt = select(existsData().where(*self._where))
        result = await self.session.scalar(stmt)
        return bool(result)

    async def not_exists(self) -> bool:
        """
        检查是否不存在满足条件的记录
        :return: True/False
        """
        return not await self.exists()

    async def delete(self) -> bool:
        """
        删除数据(物理删除)
        :return: True/False
        """
        stmt = deleteData(self.model).where(*self._where)
        result = await self.session.execute(stmt)
        await self.session.commit()
        return bool(result.rowcount)

    async def soft_delete(self, deleted_field: str = "is_deleted") -> bool:
        """
        删除数据(软删除)
        :param deleted_field: 软删除标记字段
        :return: True/False
        """
        if not hasattr(self.model, deleted_field):
            raise AttributeError(f"Model {self.model.__name__} has no '{deleted_field}' attribute")

        # 构建更新语句
        stmt = updateData(self.model).where(*self._where).values({deleted_field: True})
        result = await self.session.execute(stmt)
        await self.session.commit()
        return bool(result.rowcount)

    async def get_by_id(self, pk: Any) -> Optional[ModelT | dict[str, Any]]:
        """
        主键查单条
        :param pk: 主键
        :return:
        """
        if self._cols:
            stmt = select(self.model).options(load_only(*self._cols))
        else:
            stmt = select(self.model)
        stmt = stmt.where(self.model.id == pk).limit(1)  # type: ignore
        result = await self.session.scalar(stmt)
        return result

    async def get_one(self):
        """
        条件查一条
        :return: 返回模型对象
        """
        stmt = (
            select(self.model)
            .where(*self._where)
            .order_by(*self._order_by)
            .limit(1)
        )
        if self._cols:
            stmt = stmt.options(load_only(*self._cols))
        return await self.session.scalar(stmt)

    async def get_all(self, as_dict: bool = False) -> Sequence[Any]:
        """
        查询全部
        :return: 返回模型对象
        """
        if as_dict and self._cols:
            stmt = select(*self._cols).where(*self._where).order_by(*self._order_by)
            rows = await self.session.execute(stmt)
            return list(rows.mappings())

        stmt = select(self.model).where(*self._where).order_by(*self._order_by)
        if self._cols:
            stmt = stmt.options(load_only(*self._cols))

        rows = await self.session.execute(stmt)
        if as_dict:
            return [self._row_to_dict(row) for row in rows.scalars()]

        return rows.scalars().all()

    async def get_list(self, column: ColumnLike) -> List[Any]:
        """
        返回指定字段的列表
        :param column: 指定字段
        :return:
        """
        stmt = select(column).where(*self._where)
        if self._order_by:
            stmt = stmt.order_by(*self._order_by)

        result = await self.session.scalars(stmt)
        return list(result.all())

    async def pagination(self, page: int, page_size: int) -> PageResult:
        """
        分页查询
        :param page: 当前页
        :param page_size: 每页显示条数
        :return:
        """
        offset = (page - 1) * page_size
        #  先查总数
        count_stmt = select(func.count()).select_from(self.model).where(*self._where)
        total = (await self.session.execute(count_stmt)).scalar_one()

        #  再查当前页数据
        data_stmt = (
            select(self.model)
            .where(*self._where)
            .order_by(*self._order_by)
            .offset(offset)
            .limit(page_size)
        )
        if self._cols:
            data_stmt = data_stmt.options(load_only(*self._cols))
        items = (await self.session.execute(data_stmt)).scalars().all()

        # 计算总页数
        total_pages = (total + page_size - 1) // page_size

        return PageResult(
            items=items,
            total=total,
            total_pages=total_pages,
            page=page,
            page_size=page_size,
        )

    def _row_to_dict(self, inst: ModelT) -> dict[str, Any]:
        """ORM 实例 -> 字典（仅导出已加载的列，避免 DetachedInstanceError）"""
        if self._cols:  # 只导出指定列
            return {c: getattr(inst, c) for c in self._cols}
        # 全字段
        return {col.name: getattr(inst, col.name)
                for col in inst.__table__.columns}
